from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage
from fastapi import APIRouter
#星火认知大模型Spark3.5 Max的URL值，其他版本大模型URL值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.5/chat'
#星火认知大模型调用秘钥信息，请前往讯飞开放平台控制台（https://console.xfyun.cn/services/bm35）查看
SPARKAI_APP_ID = 'f18f8d0c000'
SPARKAI_API_SECRET = 'ZDYwZWU3MTZmMzdlOThjZTNiNjYyNzNj'
SPARKAI_API_KEY = 'c0e8b1154ddf5e4027bdfd78eb4fe2f9'
#星火认知大模型Spark3.5 Max的domain值，其他版本大模型domain值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_DOMAIN = 'generalv3.5'

router = APIRouter()

@router.get("/hello/spark/{query_msg}")
async def read_root(query_msg):
    print(query_msg)
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,
    )
    messages = [ChatMessage(
        role="user",
        content=query_msg
    )]
    handler = ChunkPrintHandler()
    a = spark.generate([messages], callbacks=[handler])
    print(a)


@router.get("/astream1/spark/{query_msg}")
def test_stream():
    from sparkai.core.callbacks import StdOutCallbackHandler
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        request_timeout=30, # 
        streaming=True,

    )
    messages = [ChatMessage(
        role="user",
        content="以高考为中心写一篇100字注意事项")]
    handler = ChunkPrintHandler()
    a = spark.generate([messages], callbacks=[handler])
    print(a)

@router.get("/astream2/spark/{query_msg}")
def test_starcoder2():
    from sparkai.log.logger import logger
    #logger.setLevel("debug")
    from sparkai.core.callbacks import StdOutCallbackHandler
    messages = [{'role': 'user',
                 'content': "以高考为中心写一篇100字注意事项"}]
    spark = ChatSparkLLM(
        spark_api_url="wss://xingchen-api.cn-huabei-1.xf-yun.com/v1.1/chat",
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=True,
        max_tokens= 1024,
    )
    messages = [
                ChatMessage(
                        role="user",
                        content=messages[0]['content']

    )]

    a = spark.stream(messages)
    for message in a:
        print(message)